package cn.kevinking.ai.mcp.provider.service.impl;

import cn.kevinking.ai.common.constant.HttpConsts;
import cn.kevinking.ai.common.model.ResultVO;
import cn.kevinking.ai.common.model.ServerInfo;
import cn.kevinking.ai.mcp.provider.McpProvider;
import cn.kevinking.ai.mcp.provider.heartbeat.BeatManager;
import cn.kevinking.ai.mcp.provider.heartbeat.BeatInfo;
import cn.kevinking.ai.mcp.provider.heartbeat.req.BeatReq;
import cn.kevinking.ai.mcp.provider.properties.McpProviderDiscoveryProperties;
import cn.kevinking.ai.mcp.provider.service.NamingService;
import cn.kevinking.ai.mcp.provider.utils.ApiUtils;
import cn.kevinking.ai.mcp.provider.utils.EnvUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeanUtils;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.http.MediaType;
import org.springframework.web.reactive.function.client.WebClient;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author kevinking
 */
@Slf4j
public class NamingHttpClient implements NamingService {
    private final WebClient webClient;
    private Map<String, ServerInfo> serverPool;
    private BeatManager beatManager;

    public NamingHttpClient(WebClient webClient, McpProviderDiscoveryProperties discoveryProperties) {
        this.webClient = webClient;
        int serverCount = this.initServers(discoveryProperties);
        this.beatManager = new BeatManager(serverCount, this);
    }

    @Override
    public void registerProvider(ServerInfo serverInfo, McpProvider instance) {
        ResultVO<ServerInfo> body = webClient.post()
                .uri(serverInfo.getUrl() + ApiUtils.REGISTER_URI)
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(instance)
                .retrieve()
                .bodyToMono(new ParameterizedTypeReference<ResultVO<ServerInfo>>() {
                })
                .block();
        if (body == null) {
            log.error("[MCP-PROVIDER] empty body from mcp server: {}, url: {}", serverInfo.getName(), serverInfo.getUrl());
            return;
        }
        if (HttpConsts.OK != body.code() || !body.data().isConnected()) {
            log.error("[MCP-PROVIDER] error from mcp server: {}", body.msg());
            return;
        }
        BeatInfo beatInfo = beatManager.newBeat(serverInfo, instance);
        beatManager.addBeatInfo(beatInfo);
    }

    @Override
    public ResultVO<Object> sendBeat(BeatInfo beatInfo) {
        BeatReq beatReq = new BeatReq();
        BeanUtils.copyProperties(beatInfo, beatReq);
        return webClient.post()
                .uri(beatInfo.getServerUrl() + ApiUtils.BEAT_URI)
                .bodyValue(beatReq)
                .retrieve()
                .bodyToMono(new ParameterizedTypeReference<ResultVO<Object>>() {})
                .block();
    }

    private int initServers(McpProviderDiscoveryProperties discoveryProperties) {
        // guarantee the server name is unique
        final AtomicInteger counter = new AtomicInteger(1);
        final int serverCount = discoveryProperties.getServers().size();
        Map<String, ServerInfo> serverMap = new ConcurrentHashMap<>(serverCount);
        discoveryProperties.getServers().forEach((serverName, value) -> {
            final String serverUrl = value.url();
            if (EnvUtils.isValidServerUrl(serverUrl)) {
                final String uniqueName = String.format("%s-%d", serverName, counter.getAndIncrement());
                serverMap.put(uniqueName, ServerInfo.builder().name(serverName).url(serverUrl).isConnected(false).build());
            } else {
                log.error("[MCP-PROVIDER] invalid server url: {}, name is {}", serverUrl, serverName);
            }
        });
        this.serverPool = serverMap;
        return serverCount;
    }
}
